import time
from typing import Dict, Any, List, Tuple
from pymodbus.client import ModbusTcpClient
from pymodbus.payload import BinaryPayloadDecoder
from pymodbus.constants import Endian
from pymodbus.exceptions import ModbusException
from config_loader import PlcConfig, PlcRegister
from utils.logger import get_logger

# Create logger instance
logger = get_logger("modbus_client.py+ModbusReader")

# Handle different versions of pymodbus Endian enum
try:
    # Newer versions use BIG/LITTLE (uppercase)
    ENDIAN_BIG = Endian.BIG
    ENDIAN_LITTLE = Endian.LITTLE
except AttributeError:
    # Older versions use Big/Little (capitalized)
    try:
        ENDIAN_BIG = Endian.Big
        ENDIAN_LITTLE = Endian.Little
    except AttributeError:
        # Fallback to numerical values
        ENDIAN_BIG = 0  # Big endian
        ENDIAN_LITTLE = 1  # Little endian


class ModbusReader:
    def __init__(self, plc_config: PlcConfig):
        self.config = plc_config
        self.client = ModbusTcpClient(
            host=self.config.connection.server_ip,
            port=self.config.connection.server_port,
            timeout=self.config.connection.timeout / 1000  # Convert to seconds
        )
        self.connected = False

    def connect(self) -> bool:
        """Connect to Modbus server"""
        logger.info(f"Attempting to connect to Modbus server at {self.config.connection.server_ip}:{self.config.connection.server_port}")
        for attempt in range(self.config.connection.retry_count):
            try:
                if self.client.connect():
                    self.connected = True
                    logger.info("Modbus connection established successfully")
                    return True
                else:
                    logger.warning(f"Modbus connection attempt {attempt + 1} failed")
                time.sleep(self.config.connection.retry_delay / 1000)
            except Exception as e:
                logger.error(f"Modbus connection failed on attempt {attempt + 1}: {str(e)}")
                time.sleep(self.config.connection.retry_delay / 1000)
        self.connected = False
        logger.error("Failed to establish Modbus connection after all retry attempts")
        return False

    def disconnect(self):
        """Disconnect"""
        if self.connected:
            self.client.close()
            self.connected = False
            logger.info("Modbus client disconnected")

    def _decode_register(self, register: PlcRegister, data: list) -> Any:
        """Decode different types of register data"""
        # Check if data is valid
        if not data or not isinstance(data, list):
            logger.error(f"Invalid data for register {register.name_en}: {data}")
            return None
            
        # Print raw data for debugging
        logger.debug(f"Raw data for {register.name_en}: {data}")
        
        # Store the last error for reporting if all attempts fail
        last_error = None
        
        # Try different combinations of byteorder and wordorder
        for byteorder in [ENDIAN_BIG, ENDIAN_LITTLE]:
            for wordorder in [ENDIAN_BIG, ENDIAN_LITTLE]:
                try:
                    decoder = BinaryPayloadDecoder.fromRegisters(
                        data,
                        byteorder=byteorder,
                        wordorder=wordorder
                    )
                    
                    if register.type == "float":
                        value = decoder.decode_32bit_float()
                        logger.debug(f"Decoded {register.name_en} as float: {value} (byteorder: {byteorder}, wordorder: {wordorder})")
                        return value
                    elif register.type == "uint16_t":
                        value = decoder.decode_16bit_uint()
                        logger.debug(f"Decoded {register.name_en} as uint16_t: {value} (byteorder: {byteorder}, wordorder: {wordorder})")
                        return value
                    elif register.type == "uint32_t":
                        value = decoder.decode_32bit_uint()
                        logger.debug(f"Decoded {register.name_en} as uint32_t: {value} (byteorder: {byteorder}, wordorder: {wordorder})")
                        return value
                    elif register.type == "uchar":
                        # Convert to ASCII string
                        try:
                            # Use BinaryPayloadDecoder to correctly extract bytes
                            # Calculate the number of bytes we want (based on register length)
                            desired_length = int(register.len) * 2 if isinstance(register.len, str) else register.len * 2
                            
                            # Extract all bytes from the decoder
                            bytes_data = decoder.decode_string(desired_length)
                            
                            # Handle the case where each character is preceded by a null byte
                            # Check if we have a null byte before each character
                            cleaned_bytes = bytes_data
                            if len(bytes_data) >= 2 and bytes_data[0] == 0:
                                # Extract every other byte starting from index 1 (the actual characters)
                                cleaned_bytes = bytes([bytes_data[i] for i in range(1, len(bytes_data), 2) if bytes_data[i] != 0])
                            
                            # Try to decode as UTF-8 first, then fall back to ASCII
                            try:
                                value = cleaned_bytes.decode('utf-8').strip('\x00').strip()
                            except UnicodeDecodeError:
                                value = cleaned_bytes.decode('ascii').strip('\x00').strip()
                            
                            logger.debug(f"Decoded {register.name_en} as uchar: '{value}' (byteorder: {byteorder}, wordorder: {wordorder})")
                            return value
                        except Exception as e:
                            logger.error(f"Error decoding uchar data for {register.name_en}: {str(e)}")
                            # Fallback: try to manually decode the data
                            try:
                                # Manually extract bytes from registers
                                bytes_data = b""
                                for reg_val in data:
                                    # Handle both int and other types
                                    if isinstance(reg_val, int):
                                        high_byte = (reg_val >> 8) & 0xFF
                                        low_byte = reg_val & 0xFF
                                        bytes_data += bytes([high_byte, low_byte])
                                    else:
                                        # Try to convert to int first
                                        int_val = int(reg_val)
                                        high_byte = (int_val >> 8) & 0xFF
                                        low_byte = int_val & 0xFF
                                        bytes_data += bytes([high_byte, low_byte])
                                
                                # Handle the case where each character is preceded by a null byte
                                cleaned_bytes = bytes_data
                                if len(bytes_data) >= 2 and bytes_data[0] == 0:
                                    # Extract every other byte starting from index 1 (the actual characters)
                                    cleaned_bytes = bytes([bytes_data[i] for i in range(1, len(bytes_data), 2) if bytes_data[i] != 0])
                                
                                # Try to decode the manually constructed bytes
                                try:
                                    value = cleaned_bytes.decode('utf-8').strip('\x00').strip()
                                except UnicodeDecodeError:
                                    value = cleaned_bytes.decode('ascii').strip('\x00').strip()
                                
                                logger.debug(f"Decoded {register.name_en} as uchar (manual method): '{value}'")
                                return value
                            except Exception as e2:
                                logger.error(f"Error in fallback decoding for {register.name_en}: {str(e2)}")
                                # Return raw data if all decoding attempts fail
                                return data
                    else:
                        logger.warning(f"Unsupported register type {register.type} for {register.name_en}")
                        return None
                        
                except AttributeError as e:
                    last_error = f"AttributeError while decoding {register.name_en} with byteorder:{byteorder}, wordorder:{wordorder}: {str(e)}"
                    logger.debug(last_error)
                    # Try to recover by converting data elements to int if they are not already
                    try:
                        int_data = [int(x) if not isinstance(x, int) else x for x in data]
                        decoder = BinaryPayloadDecoder.fromRegisters(
                            int_data,
                            byteorder=byteorder,
                            wordorder=wordorder
                        )
                        
                        if register.type == "float":
                            value = decoder.decode_32bit_float()
                            logger.debug(f"Decoded {register.name_en} as float: {value} (byteorder: {byteorder}, wordorder: {wordorder}) [after int conversion]")
                            return value
                        elif register.type == "uint16_t":
                            value = decoder.decode_16bit_uint()
                            logger.debug(f"Decoded {register.name_en} as uint16_t: {value} (byteorder: {byteorder}, wordorder: {wordorder}) [after int conversion]")
                            return value
                        elif register.type == "uint32_t":
                            value = decoder.decode_32bit_uint()
                            logger.debug(f"Decoded {register.name_en} as uint32_t: {value} (byteorder: {byteorder}, wordorder: {wordorder}) [after int conversion]")
                            return value
                    except Exception as recovery_error:
                        last_error = f"Failed to recover from AttributeError while decoding {register.name_en}: {str(recovery_error)}"
                        logger.debug(last_error)
                        continue  # Try next combination
                except Exception as e:
                    last_error = f"Failed to decode {register.name_en} with byteorder:{byteorder}, wordorder:{wordorder}: {str(e)}"
                    logger.debug(last_error)
                    continue  # Try next combination
        
        # If all combinations failed, log the last error and return raw data
        if last_error:
            logger.error(f"All decoding attempts failed for {register.name_en}. Last error: {last_error}")
        else:
            logger.error(f"All decoding attempts failed for {register.name_en}, returning raw data")
        return data

    def read_register(self, register: PlcRegister) -> Dict[str, Any]:
        """Read single register data"""
        result = {}
        try:
            # Check connection first
            if not self.connected:
                logger.info(f"Modbus not connected, attempting to connect...")
                if not self.connect():
                    logger.error(f"Cannot read register {register.name_en} due to connection failure")
                    return result
                
            for addr in register.addr:
                # Function code 03: Read holding registers
                if register.reg == "03":
                    count = int(register.len) if isinstance(register.len, str) else register.len
                    logger.debug(f"Reading register {register.name_en} at address {addr} with count {count}")
                    try:
                        response = self.client.read_holding_registers(
                            address=addr,
                            count=count,
                            slave=self.config.connection.slave_id
                        )
                    except Exception as e:
                        logger.error(f"Exception while reading register {register.name_en} at address {addr}: {str(e)}")
                        logger.error(f"Exception type: {type(e)}")
                        import traceback
                        logger.error(traceback.format_exc())
                        self.connected = False
                        return result
                    
                    logger.debug(f"Raw response for {register.name_en}: {response}")
                    if not response.isError():
                        logger.debug(f"Response registers for {register.name_en}: {response.registers}")
                        value = self._decode_register(register, response.registers)
                        # Handle multiple address cases (such as fault codes)
                        if len(register.addr) > 1:
                            if register.name_en not in result:
                                result[register.name_en] = []
                            result[register.name_en].append(value)
                        else:
                            result[register.name_en] = value
                        logger.debug(f"Successfully read register {register.name_en}: {value}")
                    else:
                        logger.error(f"Modbus read error ({register.name_en}): {response}")
                        logger.error(f"Error type: {type(response)}")
                        self.connected = False  # Mark connection exception, trigger reconnection
                time.sleep(self.config.connection.cycle_delay / 1000)
        except ModbusException as e:
            logger.error(f"Modbus read error ({register.name_en}): {str(e)}")
            logger.error(f"ModbusException type: {type(e)}")
            import traceback
            logger.error(traceback.format_exc())
            self.connected = False  # Mark connection exception, trigger reconnection
        except Exception as e:
            logger.error(f"Unexpected error reading register ({register.name_en}): {str(e)}")
            logger.error(f"Exception type: {type(e)}")
            import traceback
            logger.error(traceback.format_exc())
            self.connected = False  # Mark connection exception, trigger reconnection
        return result

    def read_all_registers(self) -> Dict[str, Any]:
        """Read all configured register data"""
        if not self.connected and not self.connect():
            return {}
            
        all_data = {}
        # Read regular registers
        for reg in self.config.registers:
            reg_data = self.read_register(reg)
            all_data.update(reg_data)
            time.sleep(self.config.connection.cycle_delay / 1000)
            
        # Read parameters
        if hasattr(self.config, 'parameters') and self.config.parameters:
            for param in self.config.parameters:
                param_data = self.read_register(param)
                all_data.update(param_data)
                time.sleep(self.config.connection.cycle_delay / 1000)
        
        logger.info(f"Completed reading all registers, total {len(all_data)} registers read")
        logger.debug(f"All register data: {all_data}")
        return all_data

    def read_all_registers_optimized(self) -> Tuple[Dict[str, Any], float]:
        """
        Read all configured register data with optimization.
        Instead of reading each register separately, read data in blocks based on address ranges.
        
        Returns:
            Tuple of (data dictionary, read time in seconds)
        """
        if not self.connected and not self.connect():
            return {}, 0.0
            
        start_time = time.time()
        all_data = {}
        
        # Combine registers and parameters for optimization
        all_registers = []
        if hasattr(self.config, 'registers') and self.config.registers:
            all_registers.extend(self.config.registers)
        if hasattr(self.config, 'parameters') and self.config.parameters:
            all_registers.extend(self.config.parameters)
            
        # Group registers by function code and address ranges
        reg_groups = self._group_registers_by_address(all_registers)
        
        # Read data in groups
        for func_code, addr_ranges in reg_groups.items():
            for start_addr, end_addr, registers_in_range in addr_ranges:
                # Calculate number of registers to read
                count = end_addr - start_addr + 1
                
                # Read the block of registers
                try:
                    if func_code == "03":  # Read holding registers
                        response = self.client.read_holding_registers(
                            address=start_addr,
                            count=count,
                            slave=self.config.connection.slave_id
                        )
                    else:
                        # Skip unsupported function codes for now
                        continue
                        
                    if not response.isError():
                        # Process the response and map to individual registers
                        self._process_register_block(response.registers, start_addr, registers_in_range, all_data)
                    else:
                        logger.error(f"Modbus read error for address range {start_addr}-{end_addr}: {response}")
                        self.connected = False
                        return {}, 0.0
                        
                except ModbusException as e:
                    logger.error(f"Modbus read error for address range {start_addr}-{end_addr}: {str(e)}")
                    self.connected = False
                    return {}, 0.0
                except Exception as e:
                    logger.error(f"Unexpected error reading address range {start_addr}-{end_addr}: {str(e)}")
                    self.connected = False
                    return {}, 0.0
                    
                # Small delay between reads
                time.sleep(self.config.connection.cycle_delay / 1000)
        
        read_time = time.time() - start_time
        logger.info(f"Completed reading all registers (optimized), total {len(all_data)} registers read in {read_time:.3f}s")
        logger.debug(f"All register data: {all_data}")
        return all_data, read_time

    def _group_registers_by_address(self, registers: List[PlcRegister]) -> Dict[str, List[Tuple[int, int, List[PlcRegister]]]]:
        """
        Group registers by function code and address ranges to minimize modbus calls.
        
        Returns:
            Dictionary mapping function code to list of (start_addr, end_addr, registers_in_range) tuples
        """
        # Group registers by function code
        func_groups = {}
        for reg in registers:
            if reg.reg not in func_groups:
                func_groups[reg.reg] = []
            func_groups[reg.reg].append(reg)
            
        # For each function code group, create address ranges
        result = {}
        for func_code, regs in func_groups.items():
            # Sort registers by address
            sorted_regs = sorted(regs, key=lambda r: min(r.addr))
            
            # Group into ranges
            ranges = []
            if sorted_regs:
                current_start = min(sorted_regs[0].addr)
                current_end = max(sorted_regs[0].addr) + (int(sorted_regs[0].len) if isinstance(sorted_regs[0].len, str) else sorted_regs[0].len) - 1
                current_regs = [sorted_regs[0]]
                
                for reg in sorted_regs[1:]:
                    reg_start = min(reg.addr)
                    reg_end = max(reg.addr) + (int(reg.len) if isinstance(reg.len, str) else reg.len) - 1
                    
                    # If this register is close to the current range, merge it
                    # We add a small buffer to avoid too many small reads
                    if reg_start <= current_end + 10:  # Within 10 registers, merge
                        current_end = max(current_end, reg_end)
                        current_regs.append(reg)
                    else:
                        # Save current range and start new one
                        ranges.append((current_start, current_end, current_regs))
                        current_start = reg_start
                        current_end = reg_end
                        current_regs = [reg]
                        
                # Add the last range
                ranges.append((current_start, current_end, current_regs))
                
            result[func_code] = ranges
            
        return result

    def _process_register_block(self, data: List[int], start_addr: int, registers: List[PlcRegister], result: Dict[str, Any]):
        """
        Process a block of register data and map to individual registers.
        
        Args:
            data: Raw register data
            start_addr: Starting address of the data block
            registers: Registers that should be extracted from this block
            result: Dictionary to store the extracted register values
        """
        for register in registers:
            # Calculate offset within the block
            reg_start_addr = min(register.addr)
            offset = reg_start_addr - start_addr
            
            if offset >= 0 and offset < len(data):
                # Extract the relevant data for this register
                reg_len = int(register.len) if isinstance(register.len, str) else register.len
                if offset + reg_len <= len(data):
                    reg_data = data[offset:offset + reg_len]
                    # Decode the register value
                    value = self._decode_register(register, reg_data)
                    result[register.name_en] = value
                    logger.debug(f"Decoded register {register.name_en}: {value}")

    def write_register(self, register: PlcRegister, value: Any) -> bool:
        """
        Write a single register value to PLC.
        
        :param register: Register to write
        :param value: Value to write
        :return: True if successful, False otherwise
        """
        try:
            # Check connection first
            if not self.connected:
                logger.info(f"Modbus not connected, attempting to connect...")
                if not self.connect():
                    logger.error(f"Cannot write register {register.name_en} due to connection failure")
                    return False
                    
            # Handle different data types
            if register.type == "uint16_t":
                if isinstance(value, (int, float)):
                    write_value = int(value)
                    address = register.addr[0]  # Use first address
                    response = self.client.write_register(
                        address=address,
                        value=write_value,
                        slave=self.config.connection.slave_id
                    )
                    
                    if not response.isError():
                        logger.debug(f"Successfully wrote {write_value} to register {register.name_en} at address {address}")
                        return True
                    else:
                        logger.error(f"Modbus write error for {register.name_en}: {response}")
                        self.connected = False
                        return False
                else:
                    logger.error(f"Invalid value type for uint16_t register {register.name_en}: {type(value)}")
                    return False
                    
            else:
                logger.warning(f"Writing to register type {register.type} for {register.name_en} not implemented")
                return False
                
        except ModbusException as e:
            logger.error(f"Modbus write error for {register.name_en}: {str(e)}")
            self.connected = False
            return False
        except Exception as e:
            logger.error(f"Unexpected error writing register {register.name_en}: {str(e)}")
            self.connected = False
            return False

    def write_registers_from_data(self, data: Dict[str, Any]) -> None:
        """
        Write register values back to PLC based on data dictionary.
        
        :param data: Dictionary containing register names and values to write
        """
        # Create a lookup for parameters by name
        param_lookup = {}
        if hasattr(self.config, 'parameters') and self.config.parameters:
            for param in self.config.parameters:
                param_lookup[param.name_en] = param
                
        # Write each value that corresponds to a parameter
        for name, value in data.items():
            if name in param_lookup:
                param = param_lookup[name]
                # Only write specific parameters we're interested in
                if name in ["modbus_read_cycle", "mqtt_publish_cycle"]:
                    # Check bounds if defined
                    if hasattr(param, 'min') and hasattr(param, 'max'):
                        if isinstance(value, (int, float)):
                            if value < param.min:
                                value = param.min
                            elif value > param.max:
                                value = param.max
                    
                    success = self.write_register(param, value)
                    if success:
                        logger.debug(f"Successfully wrote {value} to parameter {name}")
                    else:
                        logger.error(f"Failed to write {value} to parameter {name}")
